# kmeans.py
import numpy as np
import matplotlib.pyplot as plt


def euclidean_distance(p1, p2):
    """计算两点之间的欧氏距离"""
    return np.sqrt(np.sum((p1 - p2) ** 2))


class KMeans:
    def __init__(self, k=3, max_iter=300, random_state=42):
        """
        k: 聚类的簇数
        max_iter: 最大迭代次数
        random_state: 随机种子，确保每次运行结果一致
        """
        self.k = k
        self.max_iter = max_iter
        np.random.seed(random_state)
        self.centroids = None
        self.clusters = None

    def fit(self, data):
        """
        训练 K-Means 模型
        data: (n_samples, n_features) 的 numpy 数组
        """
        n_samples, n_features = data.shape

        # 1. 随机初始化质心
        random_indices = np.random.choice(n_samples, self.k, replace=False)
        self.centroids = data[random_indices]

        # 开始迭代
        for _ in range(self.max_iter):
            # 2. 分配数据点到最近的簇
            self.clusters = self._create_clusters(data)

            # 记录旧的质心以备检查收敛
            old_centroids = self.centroids.copy()

            # 3. 更新质心
            for i in range(self.k):
                if self.clusters[i]:  # 避免空簇
                    self.centroids[i] = np.mean(self.clusters[i], axis=0)

            # 4. 检查收敛：如果质心不再变化，则停止迭代
            if self._is_converged(old_centroids, self.centroids):
                break

    def _create_clusters(self, data):
        """将每个数据点分配给最近的质心"""
        clusters = [[] for _ in range(self.k)]
        for point in data:
            distances = [euclidean_distance(point, centroid) for centroid in self.centroids]
            closest_centroid_index = np.argmin(distances)
            clusters[closest_centroid_index].append(point)
        return clusters

    def _is_converged(self, old_centroids, new_centroids):
        """检查质心是否收敛"""
        distances = [euclidean_distance(old_centroids[i], new_centroids[i]) for i in range(self.k)]
        return sum(distances) == 0

    def predict(self, data):
        """预测新数据的簇标签"""
        predictions = []
        for point in data:
            distances = [euclidean_distance(point, centroid) for centroid in self.centroids]
            predictions.append(np.argmin(distances))
        return np.array(predictions)


def create_synthetic_data(n_samples=300, n_features=2, n_clusters=4, random_state=42):
    """生成一个用于测试的二维数据集"""
    np.random.seed(random_state)
    dataset = []
    cluster_std_dev = [1.0, 1.5, 0.8, 1.2]  # 让每个簇的大小不同

    # 为每个簇生成中心点
    centers = np.random.uniform(-10, 10, (n_clusters, n_features))

    points_per_cluster = n_samples // n_clusters

    for i in range(n_clusters):
        # 围绕中心点生成高斯分布的数据
        cluster_points = np.random.randn(points_per_cluster, n_features) * cluster_std_dev[i] + centers[i]
        dataset.extend(cluster_points)

    return np.array(dataset)


if __name__ == '__main__':
    # --- 1. 数据准备 ---
    K = 4
    data = create_synthetic_data(n_samples=400, n_clusters=K, random_state=10)

    # --- 2. 训练模型 ---
    print(f"开始使用 K={K} 进行 K-Means 聚类...")
    kmeans = KMeans(k=K, max_iter=100)
    kmeans.fit(data)
    print("聚类完成。")

    # --- 3. 结果可视化 ---
    plt.figure(figsize=(10, 8))
    # 使用一组漂亮的颜色
    colors = ['#377eb8', '#ff7f00', '#4daf4a', '#f781bf', '#a65628', '#984ea3']

    for i, cluster in enumerate(kmeans.clusters):
        # 将列表转换为 numpy 数组以便于索引
        cluster_points = np.array(cluster)
        if cluster_points.any():
            plt.scatter(cluster_points[:, 0], cluster_points[:, 1], c=colors[i], label=f'Cluster {i + 1}', s=50,
                        alpha=0.7)

    # 绘制最终的质心
    centroids = kmeans.centroids
    plt.scatter(centroids[:, 0], centroids[:, 1], s=250, c='black', marker='*', label='Centroids', edgecolors='white')

    plt.title('K-Means Clustering Result', fontsize=16)
    plt.xlabel('Feature 1', fontsize=12)
    plt.ylabel('Feature 2', fontsize=12)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.6)

    # 保存结果图
    plt.savefig("kmeans_result.png")
    print("\n聚类结果图已保存为 'kmeans_result.png'")

    print("\n最终的质心坐标:")
    for i, centroid in enumerate(kmeans.centroids):
        print(f"  Cluster {i + 1}: {np.round(centroid, 2)}")

    plt.show()